#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <crts.h>

#include "pcg.h"
#include "my_slave.h"
#include <algorithm>

void process_lower(const LduMatrix& ldu_matrix, CsrMatrix& csr_matrix) {
    csr_matrix.rows = ldu_matrix.cells;
    // 对角线元素单独存储
    csr_matrix.data_size = ldu_matrix.faces;
    csr_matrix.row_off = (int *)malloc((csr_matrix.rows + 2)*sizeof(int));
    csr_matrix.cols = (int *)malloc(csr_matrix.data_size*sizeof(int));
    csr_matrix.data = (double *)malloc(csr_matrix.data_size*sizeof(double));

    memset(csr_matrix.row_off, 0, (csr_matrix.rows + 2)*sizeof(int));

    int row, col, offset;

    for(int i = 0; i < ldu_matrix.faces; i++){
        row	= ldu_matrix.uPtr[i] + 2;
        csr_matrix.row_off[row]++;
    }

    for(int i = 1; i <= ldu_matrix.cells; i++){
        csr_matrix.row_off[i+1] += csr_matrix.row_off[i];
    }
    for(int i = 0; i < ldu_matrix.faces; i++){
        row = ldu_matrix.uPtr[i] + 1;
        col = ldu_matrix.lPtr[i];
        offset = csr_matrix.row_off[row]++;
        csr_matrix.cols[offset] = col;
        csr_matrix.data[offset] = ldu_matrix.lower[i];
    }
}

PCGReturn pcg_solve(const LduMatrix &ldu_matrix, double *source, double *x, int maxIter, double tolerance, double normfactor) {
    static int isInit = 0;
    if(isInit == 0){
        //从核初始化
        CRTS_init();
        isInit = 1;
    }

    // 加载输入参数 1
    LoadInputParam1 param1;
    param1.diag = ldu_matrix.diag;
    param1.source = source;
    param1.x = x;
    param1.lPtr = ldu_matrix.lPtr;
    param1.uPtr = ldu_matrix.uPtr;
    param1.data = ldu_matrix.lower;
    param1.cells = ldu_matrix.cells;
    param1.faces = ldu_matrix.faces;
    athread_spawn(SLAVE_FUN(load_input_param1), &param1);
    
    CsrMatrix csr_matrix;
    process_lower(ldu_matrix, csr_matrix);
    athread_join();

    // athread_spawn(SLAVE_FUN(load_input_param2), &csr_matrix);
    // athread_join();
    
    RunParam run_param;
    PCGReturn result;
    run_param.maxIter = maxIter;
    run_param.tolerance = tolerance;
    run_param.normfactor = normfactor;
    run_param.result = &result;
    run_param.csr_matrix = &csr_matrix;

    athread_spawn(SLAVE_FUN(run), &run_param);
    athread_join();
    free_csr_matrix(csr_matrix);
    INFO("PCG: init residual = %e, final residual = %e, iterations: %d\n", result.init_residual, result.residual, result.iter);
    return result;
}

void ldu_to_csr(const LduMatrix &ldu_matrix, CsrMatrix &csr_matrix) {
    csr_matrix.rows = ldu_matrix.cells;
    // 对角线元素单独存储
    csr_matrix.data_size = 2 * ldu_matrix.faces;
    csr_matrix.row_off = (int *)malloc((csr_matrix.rows + 1)*sizeof(int));
    csr_matrix.cols = (int *)malloc(csr_matrix.data_size*sizeof(int));
    csr_matrix.data = (double *)malloc(csr_matrix.data_size*sizeof(double));

    int row, col, offset;
    int *tmp = (int *)malloc((csr_matrix.rows + 1)*sizeof(int));

    memset(csr_matrix.row_off, 0, (csr_matrix.rows + 1)*sizeof(int));
        
    for(int i = 0; i < ldu_matrix.faces; i++){
        row	= ldu_matrix.uPtr[i] ;
        col = ldu_matrix.lPtr[i] ;
        csr_matrix.row_off[row+1]++;
        csr_matrix.row_off[col+1]++;
    }

    for(int i = 0;i< ldu_matrix.cells; i++){
        csr_matrix.row_off[i+1] += csr_matrix.row_off[i];
    }

    memcpy(&tmp[0], &csr_matrix.row_off[0], (ldu_matrix.cells + 1)*sizeof(int));
    // lower
    for(int i = 0; i < ldu_matrix.faces; i++ ){
        row = ldu_matrix.uPtr[i];
        col = ldu_matrix.lPtr[i];
        offset = tmp[row]++;
        csr_matrix.cols[offset] = col;
        csr_matrix.data[offset] = ldu_matrix.lower[i];
    }

    // upper
    for(int i = 0; i < ldu_matrix.faces; i++){
        row = ldu_matrix.lPtr[i];
        col = ldu_matrix.uPtr[i];
        offset = tmp[row]++;
        csr_matrix.cols[offset] = col;
        csr_matrix.data[offset] = ldu_matrix.upper[i];
    }

    free(tmp);
}

void csr_spmv(const CsrMatrix &csr_matrix, double *vec, double *result) {
    for(int i = 0; i < csr_matrix.rows; i++) {
        int start = csr_matrix.row_off[i];
        int num = csr_matrix.row_off[i+1] - csr_matrix.row_off[i];
        double temp = 0;
        for(int j = 0; j < num; j++) {                      
            temp += vec[csr_matrix.cols[start+j]] * csr_matrix.data[start+j]; 
        }
        result[i]=temp;
    }
}

void csr_precondition_spmv(const CsrMatrix &csr_matrix, double *vec, double *val, double *result) {
    for(int i = 0; i < csr_matrix.rows; i++) {
        int start = csr_matrix.row_off[i];
        int num = csr_matrix.row_off[i+1] - csr_matrix.row_off[i];
        double temp = 0;
        for(int j = 0; j < num; j++) {                      
            temp += vec[csr_matrix.cols[start+j]] * val[start+j]; 
        }
        result[i]=temp;
    }
}

void v_dot_product(const int nCells, const double *vec1, const double *vec2, double *result) {
    for(int cell = 0; cell < nCells; cell++) {
        result[cell] = vec1[cell] * vec2[cell];
    }
}

void v_sub_dot_product(const int nCells, const double *sub, const double *subed, const double *vec, double *result) {
    for(int cell = 0; cell < nCells; cell++) {
        result[cell] = (sub[cell] - subed[cell])*vec[cell];
    }
}

void pcg_init_precondition_csr (const CsrMatrix &csr_matrix, Precondition &pre) {
    for(int i = 0 ; i < csr_matrix.rows; i++) {
        for(int j = csr_matrix.row_off[i]; j < csr_matrix.row_off[i+1]; j++){
            if(csr_matrix.cols[j] == i) {
                pre.pre_mat_val[j] = 0.;	 
                pre.preD[i] = 1.0/csr_matrix.data[j];
            } else {
                pre.pre_mat_val[j] = csr_matrix.data[j];
            }
        }
    }
}

void pcg_precondition_csr(const CsrMatrix &csr_matrix, const Precondition &pre, double *rAPtr, double *wAPtr) {
    double* gAPtr = (double*)malloc(csr_matrix.rows*sizeof(double));
    v_dot_product(csr_matrix.rows, pre.preD, rAPtr, wAPtr);
    memset(gAPtr, 0, csr_matrix.rows*sizeof(double));
    for(int deg = 1; deg < 2; deg++) {   
        csr_precondition_spmv(csr_matrix, wAPtr, pre.pre_mat_val, gAPtr);
        v_sub_dot_product(csr_matrix.rows, rAPtr, gAPtr, pre.preD, wAPtr);
        memset(gAPtr, 0, csr_matrix.rows*sizeof(double));
    }
    free(gAPtr);
}

double pcg_gsumMag(double *r, int size) {
    double ret = .0;
    for(int i = 0; i < size; i++) {
        ret += fabs(r[i]);
    }
    return ret;
}

double pcg_gsumProd(double *z, double *r, int size) {
    double ret = .0;
    for(int i = 0; i < size; i++) {
        ret += z[i] * r[i];
    }
    return ret;
}

void free_pcg(PCG &pcg) {
    free(pcg.r);
    free(pcg.z);
    free(pcg.p);
    free(pcg.Ax);
}

void free_csr_matrix(CsrMatrix &csr_matrix) {
    free(csr_matrix.cols);
    free(csr_matrix.data);
    free(csr_matrix.row_off);
}

void free_precondition(Precondition &pre) {
    free(pre.preD);
    free(pre.pre_mat_val);
}
